package cn.edu.njust;

import cn.edu.njust.balancer.DefaultLoadBalancer;
import cn.edu.njust.balancer.LoadBalancer;
import cn.edu.njust.server.BoundedRandomLatencyServer;
import cn.edu.njust.server.Server;
import cn.edu.njust.strategy.*;
import org.junit.BeforeClass;
import org.junit.Test;

import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

public class Main {
    private static final int ROUNDS = 32;
    private static LoadBalancer balancer;
    private static List<Server> servers;
    private static List<Integer> weights;
    private static ConcurrentHashMap<String, Integer> stats;
    private static Map<String, ExecutorService> executors;

    @BeforeClass
    public static void setUp() {
        servers = List.of(
                new BoundedRandomLatencyServer("Node1", 100),
                new BoundedRandomLatencyServer("Node2", 100),
                new BoundedRandomLatencyServer("Node3", 100),
                new BoundedRandomLatencyServer("Node4", 100),
                new BoundedRandomLatencyServer("Node5", 100),
                new BoundedRandomLatencyServer("Node6", 100),
                new BoundedRandomLatencyServer("Node7", 100),
                new BoundedRandomLatencyServer("Node8", 100)
        );
        weights = List.of(
                Integer.valueOf(1),
                Integer.valueOf(2),
                Integer.valueOf(3),
                Integer.valueOf(4),
                Integer.valueOf(5),
                Integer.valueOf(6),
                Integer.valueOf(7),
                Integer.valueOf(8)
        );
        balancer = DefaultLoadBalancer
                .builder()
                .servers(servers)
                .weights(weights)
                .build();
        stats = new ConcurrentHashMap<>();
        executors = Map.of(
                "Node1", Executors.newSingleThreadExecutor(),
                "Node2", Executors.newSingleThreadExecutor(),
                "Node3", Executors.newSingleThreadExecutor(),
                "Node4", Executors.newSingleThreadExecutor(),
                "Node5", Executors.newSingleThreadExecutor(),
                "Node6", Executors.newSingleThreadExecutor(),
                "Node7", Executors.newSingleThreadExecutor(),
                "Node8", Executors.newSingleThreadExecutor()
        );
    }

    public void runTestSingleThread(String name) {
        for (int i = 0; i < ROUNDS; i++) {
            Server server = balancer.next();
            server.service();
            stats.merge(server.getId(), 1, (value, increment) -> value + increment);
            System.out.println(server);
        }
        System.out.println("Stats: " + name);
        System.out.println("============STATS=============");
        servers.forEach(node -> System.out.println(node.getId() + ": " + stats.get(node.getId())));
        System.out.println("============STATS=============");
        stats.clear();
    }

    public void runTestMultiThread(String name) throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(ROUNDS);
        for (int i = 0; i < ROUNDS; i++) {
            Server server = balancer.next();
            ExecutorService executor = executors.get(server.getId());
            executor.execute(() -> {
                server.service();
                latch.countDown();
            });
            stats.merge(server.getId(), 1, (value, increment) -> value + increment);
            System.out.println(server);
            TimeUnit.MILLISECONDS.sleep(20);
        }
        latch.await();
        System.out.println("Stats: " + name);
        System.out.println("============STATS=============");
        servers.forEach(node -> System.out.println(node.getId() + ": " + stats.get(node.getId())));
        System.out.println("============STATS=============");
        stats.clear();
    }

    @Test
    public void testRandomStrategy() {
        balancer.setStrategy(new RandomStrategy());
        balancer.init();
        runTestSingleThread("RandomStrategy");
    }

    @Test
    public void testRoundRobinStrategy() {
        balancer.setStrategy(new RoundRobinStrategy());
        balancer.init();
        runTestSingleThread("RoundRobinStrategy");
    }

    @Test
    public void testWeightedRandomStrategy() {
        balancer.setStrategy(new WeightedRandomStrategy());
        balancer.init();
        runTestSingleThread("WeightedRandomStrategy");
    }

    @Test
    public void testWeightedRoundRobinStrategy() {
        balancer.setStrategy(new WeightedRoundRobinStrategy());
        balancer.init();
        runTestSingleThread("WeightedRoundRobinStrategy");
    }

    @Test
    public void testLeastLatencyStrategy() throws InterruptedException {
        balancer.setStrategy(new LeastLatencyStrategy());
        balancer.init();
        runTestMultiThread("LeastLatencyStrategy");
    }

    @Test
    public void testLeastBacklogStrategy() throws InterruptedException {
        balancer.setStrategy(new LeastBacklogStrategy());
        balancer.init();
        runTestMultiThread("LeastBacklogStrategy");
    }
}
